import json
import zmq
import numpy as np
import cv2
import open3d as o3d
from io import BytesIO
from PIL import Image


class ZMQServer:
    """ZMQ REP server to receive color+depth images and reply best grasp.

    Usage:
        server = ZMQServer(bind_addr='tcp://*:5555', anygrasp=anygrasp, cfgs=cfgs)
        server.serve_forever()

    The server expects a multipart REQ with frames:
        0: header JSON (bytes)
        1: color image bytes (PNG/JPEG)
        2: depth image bytes (PNG) (uint16 or uint8 single-channel)

    Reply: JSON string containing either {'ok': False} or
        {'ok': True, 'translation': [x,y,z], 'rotation': [[..],[..],[..]], 'gripper_width': w, 'score': s}
    """

    def __init__(self, bind_addr='tcp://*:5555', anygrasp=None, cfgs=None):
        self.bind_addr = bind_addr
        self.anygrasp = anygrasp
        self.cfgs = cfgs

        self.ctx = zmq.Context()
        self.socket = self.ctx.socket(zmq.REP)
        self.socket.bind(self.bind_addr)

    def _decode_image(self, img_bytes):
        # decode with OpenCV from bytes
        arr = np.frombuffer(img_bytes, dtype=np.uint8)
        img = cv2.imdecode(arr, cv2.IMREAD_UNCHANGED)
        return img

    def get_obs(self):
        """Block until a request arrives, decode and return (header, color, depths).

        header: dict parsed from first frame (if any)
        color: HxWx3 RGB float32 in [0,1]
        depths: HxW array (uint16/uint8/float) matching incoming depth
        """
        while True:
            try:
                print('[ZMQServer] waiting for request...')
                frames = self.socket.recv_multipart()
            except Exception as e:
                print('ZMQ recv error in get_obs:', e)
                continue

            print(f'[ZMQServer] received request with {len(frames)} frames')

            # store frames so that reply() can send reply corresponding to this request
            self._last_frames = frames

            try:
                header = json.loads(frames[0].decode('utf-8')) if len(frames) > 0 else {}
            except Exception:
                header = {}

            # debug: print header summary
            try:
                print(f'[ZMQServer] header: {header}')
            except Exception:
                pass

            if len(frames) < 3:
                # send immediate error reply and continue waiting
                try:
                    self.socket.send_string(json.dumps({'ok': False, 'reason': 'need 3 frames'}))
                except Exception:
                    pass
                continue

            color_b = frames[1]
            depth_b = frames[2]

            try:
                print(f'[ZMQServer] color bytes={len(color_b)}, depth bytes={len(depth_b)}')
            except Exception:
                pass

            color = self._decode_image(color_b)
            depth = self._decode_image(depth_b)

            # Convert color BGR->RGB float [H,W,3]
            if color is None:
                try:
                    self.socket.send_string(json.dumps({'ok': False, 'reason': 'color decode failed'}))
                except Exception:
                    pass
                continue
            if len(color.shape) == 2:
                color = cv2.cvtColor(color, cv2.COLOR_GRAY2BGR)
            color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0

            if depth is None:
                try:
                    self.socket.send_string(json.dumps({'ok': False, 'reason': 'depth decode failed'}))
                except Exception:
                    pass
                continue

            if len(depth.shape) == 3:
                depth = depth[:, :, 0]

            depths = depth

            return header, color, depths

    def reply(self, resp_dict):
        """Send a JSON reply for the most recent request."""
        try:
            s = json.dumps(resp_dict)
            print(f'[ZMQServer] sending reply: {resp_dict}')
            self.socket.send_string(s)
        except Exception as e:
            print('ZMQ send error in reply():', e)